import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.translate.NoBatchifyTranslator;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;

import java.nio.file.Path;
import java.nio.file.Paths;

public class Test2 {
    public static void main(String[] args) {
//            [[101, 872, 1962, 2097, 1967, 1377, 809, 4500, 1408, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
//佟她孔妆叮令
        //首先创建一个模型
//        Model model = Model.newInstance("test");
        try {
            Path modeldir = Paths.get("D:\\multi_model.pt");
            NDManager manager = NDManager.newBaseManager();
            //            model.load(Paths.get("D:\\multi_model.pt"));
//            System.out.println(model);

            Translator<String, String> translator = new Translator<String, String>() {
                @Override
                public NDList processInput(TranslatorContext ctx, String input) throws Exception {
//                    NDManager ndManager = ctx.getNDManager();
//                    NDArray ndArray = ndManager.create(input);
                    NDArray nd = manager.create(new int[]{101, 872, 1962, 2097, 1967, 1377, 809, 4500, 1408, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
                    System.out.println(nd);
                    return new NDList(nd);
                }

                @Override
                public String processOutput(TranslatorContext ctx, NDList list) throws Exception {
//                    System.out.println(list.get(0));
                    return list.get(0).toDebugString();
                }
            };

            // 加载模型
            Criteria<String, String> criteria = Criteria.builder()
                    .setTypes(String.class, String.class)
                    .optModelPath(modeldir)
                    .optTranslator(translator)

                    .build();
            // 之后调用predictor，生成预测器
//            Predictor<NDList, NDList> predictor = criteria.loadModel().newPredictor();
            Predictor<String, String> predictor = ModelZoo.loadModel(criteria).newPredictor();

            System.out.println("=========");

//            int[] inLongArray = new int[]{101, 872, 1962, 2097, 1967, 1377, 809, 4500, 1408, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

//            NDList predict = predictor.predict(ndArrays);
            String predict = predictor.predict("[[101, 872, 1962, 2097, 1967, 1377, 809, 4500, 1408, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]");
            System.out.println("predict: " + predict);
        } catch (Exception e) {
            System.out.println("11 ");
            e.printStackTrace();
        }
    }
}